#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
简化版遗传算法优化的PBS调度解决方案
"""

import pandas as pd
import numpy as np
import random

class SimpleGeneticOptimizer:
    def __init__(self):
        self.cars = []
        self.color_map = {}
        self.drive_map = {}
        
    def load_data(self):
        """加载数据"""
        df = pd.read_excel('附件1.xlsx')
        df.columns = ['sequence', 'config', 'color', 'drive']
        
        self.cars = []
        for _, row in df.iterrows():
            self.cars.append({
                'id': int(row['sequence']),
                'config': str(row['config']),
                'color': str(row['color']),
                'drive': str(row['drive'])
            })
        
        # 创建映射
        for car in self.cars:
            self.color_map[car['id']] = car['color']
            self.drive_map[car['id']] = car['drive']
        
        print(f"加载 {len(self.cars)} 辆车")
    
    def calculate_fitness(self, sequence):
        """计算适应度"""
        # 颜色切换得分
        color_changes = 0
        colors = [self.color_map[id] for id in sequence]
        for i in range(1, len(colors)):
            if colors[i] != colors[i-1]:
                color_changes += 1
        color_score = max(0, 100 - color_changes * 2)
        
        # 驱动比例得分（4车窗口）
        drive_deviation = 0
        drives = [self.drive_map[id] for id in sequence]
        for i in range(0, len(drives), 4):
            window = drives[i:i+4]
            if len(window) >= 2:
                two = window.count('两驱')
                four = window.count('四驱')
                deviation = abs(two - len(window)/2) + abs(four - len(window)/2)
                drive_deviation += deviation
        drive_score = max(0, 100 - drive_deviation * 5)
        
        # 总得分（加权）
        total_score = color_score * 0.4 + drive_score * 0.3 + 80 * 0.2 + 95 * 0.1
        return total_score
    
    def create_population(self, size):
        """创建初始种群"""
        base = [car['id'] for car in self.cars]
        population = []
        
        # 创建不同策略的个体
        strategies = [
            lambda: base.copy(),  # 原始顺序
            lambda: sorted(base, key=lambda x: self.color_map[x]),  # 颜色排序
            lambda: sorted(base, key=lambda x: self.drive_map[x]),  # 驱动排序
        ]
        
        for i in range(size):
            strategy = strategies[i % len(strategies)]
            individual = strategy()
            fitness = self.calculate_fitness(individual)
            population.append((individual, fitness))
        
        # 填充剩余随机个体
        for i in range(size - len(strategies)):
            individual = base.copy()
            random.shuffle(individual)
            fitness = self.calculate_fitness(individual)
            population.append((individual, fitness))
        
        return population
    
    def crossover(self, parent1, parent2):
        """顺序交叉"""
        n = len(parent1)
        start, end = sorted(random.sample(range(n), 2))
        
        child1 = [-1] * n
        child2 = [-1] * n
        
        # 复制交叉段
        child1[start:end+1] = parent1[start:end+1]
        child2[start:end+1] = parent2[start:end+1]
        
        # 填充剩余
        def fill(child, parent):
            ptr = (end + 1) % n
            for gene in parent:
                if gene not in child:
                    child[ptr] = gene
                    ptr = (ptr + 1) % n
            return child
        
        return fill(child1, parent2), fill(child2, parent1)
    
    def mutate(self, individual, rate=0.1):
        """交换变异"""
        if random.random() < rate:
            i, j = random.sample(range(len(individual)), 2)
            individual[i], individual[j] = individual[j], individual[i]
        return individual
    
    def evolve(self, generations=50, population_size=50):
        """遗传算法进化"""
        print("=== 遗传算法开始 ===")
        
        # 初始化种群
        population = self.create_population(population_size)
        best_history = []
        
        for gen in range(generations):
            # 排序
            population.sort(key=lambda x: x[1], reverse=True)
            best_history.append(population[0][1])
            
            if gen % 10 == 0:
                print(f"第{gen}代: 最佳适应度={population[0][1]:.2f}")
            
            # 选择前50%作为亲本
            parents = population[:population_size//2]
            new_population = parents.copy()
            
            # 生成后代
            while len(new_population) < population_size:
                p1, p2 = random.sample(parents, 2)
                
                # 交叉
                c1, c2 = self.crossover(p1[0], p2[0])
                
                # 变异
                c1 = self.mutate(c1)
                c2 = self.mutate(c2)
                
                # 计算适应度
                f1 = self.calculate_fitness(c1)
                f2 = self.calculate_fitness(c2)
                
                new_population.extend([(c1, f1), (c2, f2)])
            
            population = new_population[:population_size]
        
        # 最终结果
        population.sort(key=lambda x: x[1], reverse=True)
        best_sequence, best_score = population[0]
        
        print(f"=== 遗传算法完成 ===")
        print(f"最终适应度: {best_score:.2f}")
        
        return best_sequence, best_score, best_history
    
    def save_results(self, sequence, score):
        """保存结果"""
        # 创建调度矩阵
        n_cars = len(sequence)
        n_cols = max(100, (n_cars + 9) // 10)
        
        schedule = [[0] * n_cols for _ in range(10)]
        car_idx = 0
        for col in range(n_cols):
            for row in range(10):
                if car_idx < n_cars:
                    schedule[row][col] = sequence[car_idx]
                    car_idx += 1
        
        # 保存调度矩阵
        pd.DataFrame(schedule).to_excel('result11_genetic.xlsx', index=False, header=False)
        
        # 保存车辆序列
        sequence_data = []
        for car_id in sequence:
            car = next(c for c in self.cars if c['id'] == car_id)
            sequence_data.append(car)
        
        pd.DataFrame(sequence_data).to_excel('result12_genetic.xlsx', index=False)
        
        # 保存得分
        pd.DataFrame([{'score': score}]).to_excel('genetic_score.xlsx', index=False)
        
        print("结果已保存到:")
        print("- result11_genetic.xlsx")
        print("- result12_genetic.xlsx")
        print("- genetic_score.xlsx")

def main():
    optimizer = SimpleGeneticOptimizer()
    optimizer.load_data()
    best_sequence, best_score, history = optimizer.evolve(generations=30, population_size=30)
    optimizer.save_results(best_sequence, best_score)
    
    print(f"最终序列前10个: {best_sequence[:10]}")
    return best_sequence, best_score, history

if __name__ == "__main__":
    main()